import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import pandas as pd
import os
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from loss.loss_function import FocalLoss, AsymmetricLoss
from tensorflow.python.keras.losses import BinaryCrossentropy

# 假设你已经加载了训练数据
train_data = pd.read_csv(r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/train/_classes.csv')
val_data = pd.read_csv(r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/valid/_classes.csv')
test_data = pd.read_csv(r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/test/_classes.csv')

train_dir = r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/train'
val_dir = r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/valid'
test_dir = r'/home/lsh2022400251/dataset/Tongue.v1i.multiclass/test'

# 使用 tf.data.Dataset 加载图像并缓存到内存中
def create_dataset(data, directory, batch_size=32):
    # 使用 tf.data.Dataset 创建数据集
    dataset = tf.data.Dataset.from_tensor_slices((data['filename'].values, data.iloc[:, 1:].values))

    def load_and_preprocess_image(filename, label):
        # 使用 tf.strings.join 来拼接路径
        path = tf.strings.join([directory, "/", filename])  # 拼接文件路径
        img = tf.io.read_file(path)  # 读取图像文件
        img = tf.image.decode_jpeg(img, channels=3)
        img = tf.image.resize(img, [640, 640])
        img = img / 255.0  # 归一化
        return img, label

    # 使用 map 进行图像加载和预处理
    dataset = dataset.map(lambda filename, label: load_and_preprocess_image(filename, label),
                          num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.shuffle(len(data))  # 打乱数据
    dataset = dataset.batch(batch_size)  # 批量加载数据
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)  # 预取数据
    return dataset

# 创建训练集、验证集和测试集的数据集对象
train_dataset = create_dataset(train_data, train_dir, batch_size=90)
val_dataset = create_dataset(val_data, val_dir, batch_size=90)
test_dataset = create_dataset(test_data, test_dir, batch_size=90)
strategy = tf.distribute.MirroredStrategy()

# 在分布式策略下构建模型
with strategy.scope():
    # 构建 ResNet152 模型
    base_model = tf.keras.applications.ResNet152(weights='imagenet', include_top=False, input_shape=(640, 640, 3))
    # base_model.trainable = False

    model = tf.keras.Sequential([
        base_model,
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(1024, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(train_data.shape[1] - 1,activation="sigmoid")  # 多标签分类
    ])
    ASL = AsymmetricLoss() 
    # loss_function = FocalLoss()
    # loss_function = BinaryCrossentropy()
    # 编译模型
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
                  loss=ASL,
                  metrics=[
                        tf.keras.metrics.BinaryAccuracy(name='accuracy')
                      ])

model.fit(train_dataset, validation_data=val_dataset, epochs=30)

# 评估模型
score = model.evaluate(test_dataset)
print(f'测试集损失: {score[0]}')
print(f'测试集准确率: {score[1]}')

# 保存模型
model.save('image_classification_model.h5')

